import numpy as np
import argparse
import cv2
from PIL import Image
import traceback


class ImageFeature:
    """
    特征提取结果类，保存ROI图片和特征点
    """
    def __init__(self, src_img=None, roi_img=None, keypoints=None, descriptors=None, roi_rect=None):
        """
        初始化ImageFeature类
        
        参数:
        src_img: 源图像
        roi_img: ROI区域图像
        keypoints: 关键点列表
        descriptors: 描述子矩阵
        rect: 截取的矩形区域 (x, y, w, h)
        """
        self.src_img = src_img.copy()
        self.roi_img = roi_img.copy()
        self.descriptors = descriptors
        self.roi_rect = roi_rect
        self.keypoints = keypoints
        
        # 创建偏移后的关键点
        self.offset_keypoints = []
        if roi_rect is not None and keypoints is not None:
            x, y, w, h = roi_rect
            for kp in keypoints:
                # 创建新的KeyPoint对象，添加坐标偏移
                new_kp = cv2.KeyPoint(kp.pt[0] + x, kp.pt[1] + y, kp.size, kp.angle, 
                                    kp.response, kp.octave, kp.class_id)
                self.offset_keypoints.append(new_kp)
        else:
            self.offset_keypoints = keypoints.copy() if keypoints else []

class MatchResult:
    """
    匹配结果的类
    """
    def __init__(self):
        """
        初始化MatchResult类
        参数:
        """
        self.score = 0              # 匹配得分
        self.matched_points = []    # 匹配点对
        self.homography_matrix = None  # 单应性矩阵
        self.warped_img = None      # 透视变换后的图像
    

class SIFTMatcher:
    """
    SIFT特征提取和匹配工具类
    """
    
    def __init__(self, ratio_threshold=0.7, min_matches=10):
        """
        初始化SIFT匹配器
        
        参数:
        ratio_threshold: 特征匹配的比率阈值（默认：0.7）
        min_matches: 最小匹配点数阈值（默认：10）
        """
        self.ratio_threshold = ratio_threshold
        self.min_matches = min_matches
        try:
            self.sift_detector = cv2.SIFT_create()
        except Exception as e:
            print(f"初始化SIFT检测器失败: {e}")
            # 对于较新版本的OpenCV，尝试使用xfeatures2d
            try:
                self.sift_detector = cv2.xfeatures2d.SIFT_create()
            except Exception as e2:
                print(f"使用xfeatures2d也失败: {e2}")
                raise Exception("无法初始化SIFT检测器，请确保OpenCV版本支持SIFT")
    
    def feature_extract(self, img: np.ndarray, rect=None) -> ImageFeature:
        """
        提取图像的sift特征点
        
        参数:
        img: 输入图像
        rect: 感兴趣区域（ROI），格式为 (x, y, w, h)
        
        返回:
        ImageFeature对象
        """
        try:
            # 处理ROI截取
            roi_img = img
            roi_rect = None
            
            if rect is not None:
                x, y, w, h = rect
                # 确保坐标在图像范围内
                h_img, w_img = img.shape[:2]
                x = max(0, min(x, w_img - 1))
                y = max(0, min(y, h_img - 1))
                w = max(1, min(w, w_img - x))
                h = max(1, min(h, h_img - y))
                
                # 截取ROI
                roi_img = img[y:y+h, x:x+w]
                roi_rect = (x, y, w, h)
                
                print(f"截取ROI区域: ({x}, {y}, {w}, {h})")
            
            # 确保图像是灰度图像
            if len(roi_img.shape) == 3:
                gray_img = cv2.cvtColor(roi_img, cv2.COLOR_BGR2GRAY)
            else:
                gray_img = roi_img.copy()
            
            # 提取SIFT特征点和描述子
            keypoints, descriptors = self.sift_detector.detectAndCompute(gray_img, None)
            # 过滤掉质量较差的关键点
            if descriptors is not None:
                # 计算每个关键点的响应强度（用于过滤）
                responses = [kp.response for kp in keypoints]
                threshold = np.percentile(responses, 30) if responses else 0
                
                # 过滤关键点
                filtered_keypoints = []
                filtered_descriptors = []
                for kp, desc in zip(keypoints, descriptors):
                    if kp.response > threshold:
                        filtered_keypoints.append(kp)
                        filtered_descriptors.append(desc)
                
                keypoints = filtered_keypoints
                descriptors = np.array(filtered_descriptors) if filtered_descriptors else None
            # 创建ImageFeature对象
            feature_result = ImageFeature(img, roi_img, keypoints, descriptors, roi_rect)
            
            print(f"成功提取 {len(keypoints)} 个SIFT特征点")
            return feature_result
            
        except Exception as e:
            print(f"特征提取失败: {e}")
            traceback.print_exc()
            return None
    
    def feature_match(self, feature: ImageFeature, img: np.ndarray, is_plot=False) -> MatchResult:
        """
        特征匹配
        
        参数:
        feature: 待匹配的特征
        img: 待匹配的图像
        is_plot: 是否绘制匹配结果
        返回:
        MatchResult对象
        """
        try:
            # 检查输入有效性
            if feature is None or feature.descriptors is None or len(feature.descriptors) == 0:
                print("错误：提供的特征对象无效或没有描述子")
                return None
                
            if img is None:
                print("错误：目标图像为空")
                return None
            
            # 提取目标图像的特征
            target_feature = self.feature_extract(img)
            if target_feature is None or target_feature.descriptors is None:
                print("错误：无法提取目标图像特征")
                return None
            
            # 使用FLANN匹配器进行特征匹配
            FLANN_INDEX_KDTREE = 1
            index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
            search_params = dict(checks=50)
            
            flann = cv2.FlannBasedMatcher(index_params, search_params)
            
            # 执行KNN匹配
            matches = flann.knnMatch(feature.descriptors, target_feature.descriptors, k=2)
            
            # 使用更严格的Lowe's ratio test筛选好的匹配，提高匹配精度
            good_matches = []
            for m, n in matches:
                if m.distance < 0.6 * n.distance:  # 降低阈值，提高匹配精度
                    good_matches.append(m)
            
            # 创建匹配结果对象
            result = MatchResult()
            result.matched_points = good_matches
            result.score = len(good_matches)
            
            # 提取匹配点的坐标
            dst_pts = np.float32([feature.offset_keypoints[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
            src_pts = np.float32([target_feature.keypoints[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
            
            # 使用更严格的RANSAC参数计算单应性矩阵，提高精度
            M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 2.0)  # 降低容错阈值
            
            if M is not None:
                result.homography_matrix = M
                
                # 如果有ROI，使用rect起始点作为offset对单应性矩阵做偏移
                # if feature.rect is not None:
                #     x, y, w, h = feature.rect
                #     offset_matrix = np.array([[1, 0, x], [0, 1, y], [0, 0, 1]], dtype=np.float32)
                #     M = offset_matrix @ M
                
                # 执行透视变换
                h, w = img.shape[:2]
                warped_img = cv2.warpPerspective(img, M, (w, h))
                result.warped_img = warped_img
                # 如果需要绘制匹配结果
                if is_plot:
                    # 绘制匹配线 - 直接使用偏移后的关键点
                    img_matches = cv2.drawMatches(feature.src_img, feature.offset_keypoints, 
                                                img, target_feature.keypoints,
                                                good_matches, None,
                                                flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
                    
                    # 显示匹配结果
                    cv2.imshow('High Precision Feature Matching', img_matches)
                    cv2.waitKey(0)
                    cv2.destroyAllWindows()
            
            print(f"高精度匹配完成，找到 {len(good_matches)} 个匹配点")
            return result
            
        except Exception as e:
            print(f"特征匹配失败: {e}")
            traceback.print_exc()
            return None


class RectangleDrawer:
    """
    矩形框绘制测试类，用于交互式选择ROI区域
    """
    def __init__(self, image):
        """
        初始化矩形框绘制器
        
        参数:
        image: 要显示的图像
        """
        self.image = image.copy()
        self.original_image = image.copy()
        self.drawing = False
        self.rect_start = None
        self.rect_end = None
        self.rect = None
        
    def draw_rectangle_callback(self, event, x, y, flags, param):
        """
        鼠标回调函数，用于绘制矩形框
        """
        if event == cv2.EVENT_LBUTTONDOWN:
            self.drawing = True
            self.rect_start = (x, y)
            self.rect_end = (x, y)
        elif event == cv2.EVENT_MOUSEMOVE:
            if self.drawing:
                self.rect_end = (x, y)
                self.image = self.original_image.copy()
                cv2.rectangle(self.image, self.rect_start, self.rect_end, (0, 255, 0), 2)
                cv2.imshow('Image 1 - Draw Rectangle', self.image)
        elif event == cv2.EVENT_LBUTTONUP:
            self.drawing = False
            self.rect_end = (x, y)
            self.image = self.original_image.copy()
            cv2.rectangle(self.image, self.rect_start, self.rect_end, (0, 255, 0), 2)
            cv2.imshow('Image 1 - Draw Rectangle', self.image)
    
    def get_rectangle(self):
        """
        获取用户选择的矩形区域
        
        返回:
        rect: 矩形区域 (x, y, w, h) 或 None（如果用户取消）
        """
        cv2.namedWindow('Image 1 - Draw Rectangle')
        cv2.setMouseCallback('Image 1 - Draw Rectangle', self.draw_rectangle_callback)
        
        print("请在第一张图片上绘制矩形框（按ESC键取消，按回车键确认）")
        
        while True:
            cv2.imshow('Image 1 - Draw Rectangle', self.image)
            key = cv2.waitKey(1) & 0xFF
            
            if key == 27:  # ESC键
                print("用户取消操作")
                cv2.destroyAllWindows()
                return None
            elif key == 13:  # 回车键
                if self.rect_start is not None and self.rect_end is not None:
                    break
                else:
                    print("请先绘制矩形框")
        
        # 计算矩形区域
        x1, y1 = self.rect_start
        x2, y2 = self.rect_end
        
        # 确保坐标顺序正确
        x = min(x1, x2)
        y = min(y1, y2)
        w = abs(x2 - x1)
        h = abs(y2 - y1)
        
        # 确保矩形有一定大小
        if w < 10 or h < 10:
            print("矩形框太小，使用整个图像")
            self.rect = None
        else:
            self.rect = (x, y, w, h)
            print(f"用户选择的矩形区域: ({x}, {y}, {w}, {h})")
        
        cv2.destroyAllWindows()
        return self.rect


# -----------------------------
# main: 调用与测试
# -----------------------------
def main():
    """
    主函数：处理命令行参数并执行图像匹配测试
    """
    parser = argparse.ArgumentParser(description='图片SIFT特征提取和匹配工具')
    
    # 添加命令行参数
    parser.add_argument('image1', help='第一张图片路径（作为模板）')
    parser.add_argument('image2', help='第二张图片路径（待匹配）')
    parser.add_argument('--show', action='store_true', help='显示匹配结果可视化')
    parser.add_argument('--threshold', type=float, default=0.7, 
                       help='特征匹配的比率阈值（默认：0.7）')
    parser.add_argument('--min_matches', type=int, default=10,
                       help='最小匹配点数阈值（默认：10）')
    
    args = parser.parse_args()
    
    try:
        # 创建SIFT匹配器实例
        print("正在初始化SIFT匹配器...")
        matcher = SIFTMatcher(
            ratio_threshold=args.threshold,
            min_matches=args.min_matches
        )
        
        # 读取图像
        print(f"正在读取图像1: {args.image1}")
        img1 = cv2.imread(args.image1)
        if img1 is None:
            print(f"错误：无法读取图像1: {args.image1}")
            return
        
        print(f"正在读取图像2: {args.image2}")
        img2 = cv2.imread(args.image2)
        if img2 is None:
            print(f"错误：无法读取图像2: {args.image2}")
            return
        
        # 使用RectangleDrawer类让用户绘制矩形框
        drawer = RectangleDrawer(img1)
        rect = drawer.get_rectangle()
        
        if rect is None:
            print("用户取消了操作")
            return
        
        # 提取第一张图像的特征（使用用户选择的ROI）
        print("正在提取第一张图像的SIFT特征...")
        feature = matcher.feature_extract(img1, rect=rect)
        if feature is None:
            print("错误：无法提取第一张图像的特征")
            return
        
        # 执行特征匹配
        print("正在进行特征匹配...")
        result = matcher.feature_match(feature, img2, is_plot=args.show)
        
        if result is None:
            print("错误：图像匹配失败")
            return
        
        # 输出匹配结果
        print(f"\n匹配结果:")
        print(f"匹配得分: {result.score}")
        print(f"匹配点数: {len(result.matched_points)}")
        
        # 输出单应性矩阵信息
        if result.homography_matrix is None:
            print("无法计算单应性矩阵（匹配点不足）")
            return
        
        # 显示变换后的图像
        if result.warped_img is not None:
            print(f"变换完成，变换后图像尺寸: {result.warped_img.shape[:2]}")
            cv2.imwrite("warped_image.png", result.warped_img)
            
            # 如果需要显示结果
            if args.show:
                # 创建对比显示
                comparison = np.hstack((img2, result.warped_img))
                
                # 添加分割线
                cv2.line(comparison, (img2.shape[1], 0), (img2.shape[1], img2.shape[0]), (0, 255, 0), 2)
                
                # 添加文字标签
                cv2.putText(comparison, 'Original', (10, 30), 
                           cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
                cv2.putText(comparison, 'Warped', (img2.shape[1] + 10, 30), 
                           cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
                # 添加roi矩形框
                cv2.imshow('Feature Matching Result', comparison)
                print("按任意键退出...")
                cv2.waitKey(0)
                cv2.destroyAllWindows()
        else:
            print("变换失败")

            
    except Exception as e:
        print(f"程序执行出错: {e}")
        traceback.print_exc()

if __name__ == "__main__":
    main()